"""
官网地址参照：https://modelscope.cn/models/OpenBMB/MiniCPM-V-2_6/summary
利用 OCR 技术解析图片文字的
"""

import torch
from PIL import Image
from modelscope import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained('OpenBMB/MiniCPM-V-2_6', trust_remote_code=True,
    attn_implementation='sdpa', torch_dtype=torch.bfloat16) # sdpa or flash_attention_2, no eager
model = model.eval().cuda()
tokenizer = AutoTokenizer.from_pretrained('OpenBMB/MiniCPM-V-2_6', trust_remote_code=True)

# 这里是图片打开
image = Image.open('/data/resource/test.png').convert('RGB')
# 问题
question = 'What is in the image?'
# 发送的消息
msgs = [{'role': 'user', 'content': [image, question]}]

# 以交互的形式发送给模型
res = model.chat(
    image=None,
    msgs=msgs,
    tokenizer=tokenizer
)
print(res)

## if you want to use streaming, please make sure sampling=True and stream=True
## the model.chat will return a generator
# 这里是流式输出
res = model.chat(
    image=None,
    msgs=msgs,
    tokenizer=tokenizer,
    sampling=True,
    stream=True
)

generated_text = ""
for new_text in res:
    generated_text += new_text
    print(new_text, flush=True, end='')